{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "##  使用ConvLstm模型，废了很大劲修改模型，线上22.7分\n",
    "##  ConvLSTM_pytorch   (https://github.com/ndrplz/ConvLSTM_pytorch)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from copy import deepcopy\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import torch.nn as nn\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import zipfile\n",
    "import torchvision.models as models\n",
    "\n",
    "# python ref/train.py \n",
    "def set_seed(seed = 123):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    torch.manual_seed(seed)\n",
    "\n",
    "\n",
    "def load_data2():\n",
    "    # CMIP data    \n",
    "    #PATH = './tcdata/enso_round1_train_20210201/'\n",
    "    #PATH = 'E:/enso_round1_train_20210201/'\n",
    "    PATH = './data/'\n",
    "    train = xr.open_dataset(PATH+'CMIP_train.nc')\n",
    "    label = xr.open_dataset(PATH+'CMIP_label.nc')  \n",
    "   \n",
    "    train_sst = train['sst'][:, :12].values  # (4645, 12, 24, 72)\n",
    "    train_t300 = train['t300'][:, :12].values\n",
    "    train_ua = train['ua'][:, :12].values\n",
    "    train_va = train['va'][:, :12].values\n",
    "    train_label = label['nino'][:, 12:36].values\n",
    "    \n",
    "    m = np.zeros(train_ua.shape[0])\n",
    "    for i in range(train_ua.shape[0]):\n",
    "        if np.sum(np.isnan(train_ua[i])) != 0:\n",
    "            m[i] = np.sum(np.isnan(train_ua[i]))\n",
    "\n",
    "    #将有缺失值的数据 剔除，选择缺失值为0的数据\n",
    "    train_sst = train_sst[np.where(m==0)[0]]\n",
    "    train_t300 = train_t300[np.where(m==0)[0]]\n",
    "    train_ua = train_ua[np.where(m==0)[0]]\n",
    "    train_va = train_va[np.where(m==0)[0]]\n",
    "    train_label = train_label[np.where(m==0)[0]]\n",
    "\n",
    "    N = int(len(train_label)*0.8)\n",
    "    print('Train samples: {}, Valid samples: {}'.format(len(train_label[:N]), len(train_label[N:])))\n",
    "\n",
    "    dict_train = {\n",
    "        'sst':train_sst[:N],\n",
    "        't300':train_t300[:N],\n",
    "        'ua':train_ua[:N],\n",
    "        'va': train_va[:N],\n",
    "        'label': train_label[:N]}\n",
    "    dict_valid = {\n",
    "        'sst':train_sst[N:],\n",
    "        't300':train_t300[N:],\n",
    "        'ua':train_ua[N:],\n",
    "        'va': train_va[N:],\n",
    "        'label': train_label[N:]}\n",
    "    \n",
    "    train_dataset = EarthDataSet(dict_train)\n",
    "    valid_dataset = EarthDataSet(dict_valid)\n",
    "    return train_dataset, valid_dataset\n",
    "    \n",
    "\n",
    "class EarthDataSet(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data['sst'])\n",
    "\n",
    "    def __getitem__(self, idx):   \n",
    "        return (self.data['sst'][idx], self.data['t300'][idx], self.data['ua'][idx], self.data['va'][idx]), self.data['label'][idx]\n",
    "\n",
    "class ConvLSTMCell(nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, hidden_dim, kernel_size, bias):\n",
    "        \"\"\"\n",
    "        Initialize ConvLSTM cell.\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_dim: int\n",
    "            Number of channels of input tensor.\n",
    "        hidden_dim: int\n",
    "            Number of channels of hidden state.\n",
    "        kernel_size: (int, int)\n",
    "            Size of the convolutional kernel.\n",
    "        bias: bool\n",
    "            Whether or not to add the bias.\n",
    "        \"\"\"\n",
    "\n",
    "        super(ConvLSTMCell, self).__init__()\n",
    "        \n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.kernel_size = kernel_size\n",
    "        self.padding = kernel_size[0] // 2, kernel_size[1] // 2\n",
    "        self.bias = bias\n",
    "\n",
    "        self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,\n",
    "                              out_channels=4 * self.hidden_dim,\n",
    "                              kernel_size=self.kernel_size,\n",
    "                              padding=self.padding,\n",
    "                              bias=self.bias).to(device)\n",
    "\n",
    "    def forward(self, input_tensor, cur_state):\n",
    "        h_cur, c_cur = cur_state\n",
    "\n",
    "        combined = torch.cat([input_tensor, h_cur], dim=1)  # concatenate along channel axis\n",
    "\n",
    "        combined_conv = self.conv(combined)\n",
    "        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)\n",
    "        i = torch.sigmoid(cc_i)\n",
    "        f = torch.sigmoid(cc_f)\n",
    "        o = torch.sigmoid(cc_o)\n",
    "        g = torch.tanh(cc_g)\n",
    "\n",
    "        c_next = f * c_cur + i * g\n",
    "        h_next = o * torch.tanh(c_next)\n",
    "\n",
    "        return h_next.to(device), c_next.to(device)\n",
    "\n",
    "    def init_hidden(self, batch_size, image_size):\n",
    "        height, width = image_size\n",
    "        return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),\n",
    "                torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))\n",
    "\n",
    "class ConvLSTM(nn.Module):\n",
    "\n",
    "    \"\"\"\n",
    "    Parameters:\n",
    "        input_dim: Number of channels in input\n",
    "        hidden_dim: Number of hidden channels\n",
    "        kernel_size: Size of kernel in convolutions\n",
    "        num_layers: Number of LSTM layers stacked on each other\n",
    "        batch_first: Whether or not dimension 0 is the batch or not\n",
    "        bias: Bias or no bias in Convolution\n",
    "        return_all_layers: Return the list of computations for all layers\n",
    "        Note: Will do same padding.\n",
    "    Input:\n",
    "        A tensor of size B, T, C, H, W or T, B, C, H, W\n",
    "    Output:\n",
    "        A tuple of two lists of length num_layers (or length 1 if return_all_layers is False).\n",
    "            0 - layer_output_list is the list of lists of length T of each output\n",
    "            1 - last_state_list is the list of last states\n",
    "                    each element of the list is a tuple (h, c) for hidden state and memory\n",
    "    Example:\n",
    "        >> x = torch.rand((32, 10, 64, 128, 128))\n",
    "        >> convlstm = ConvLSTM(64, 16, 3, 1, True, True, False)\n",
    "        >> _, last_states = convlstm(x)\n",
    "        >> h = last_states[0][0]  # 0 for layer index, 0 for h index\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_dim, hidden_dim, kernel_size, num_layers,\n",
    "                 batch_first=False, bias=True, return_all_layers=False):\n",
    "        super(ConvLSTM, self).__init__()\n",
    "\n",
    "        #self._check_kernel_size_consistency(kernel_size)\n",
    "\n",
    "        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers\n",
    "        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)\n",
    "        hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)\n",
    "        if not len(kernel_size) == len(hidden_dim) == num_layers:\n",
    "            raise ValueError('Inconsistent list length.')\n",
    "\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.kernel_size = kernel_size\n",
    "        self.num_layers = num_layers\n",
    "        self.batch_first = batch_first\n",
    "        self.bias = bias\n",
    "        self.return_all_layers = return_all_layers\n",
    "\n",
    "        cell_list = []\n",
    "        for i in range(0, self.num_layers):\n",
    "            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]\n",
    "            cell_list.append(ConvLSTMCell(input_dim=cur_input_dim,\n",
    "                                          hidden_dim=self.hidden_dim[i],\n",
    "                                          kernel_size=self.kernel_size[i],\n",
    "                                          bias=self.bias))\n",
    "\n",
    "        self.cell_list = nn.ModuleList(cell_list)\n",
    "\n",
    "        self.liner = nn.Linear(8*24*72,24)\n",
    "    def forward(self, sst, t300, ua, va, hidden_state=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_tensor: todo\n",
    "            5-D Tensor either of shape (t, b, c, h, w) or (b, t, c, h, w)\n",
    "        hidden_state: todo\n",
    "            None. todo implement stateful\n",
    "        Returns\n",
    "        -------\n",
    "        last_state_list, layer_output\n",
    "        \"\"\"\n",
    "        \n",
    "        \n",
    "        input_tensor = torch.cat([sst.unsqueeze(2),t300.unsqueeze(2),ua.unsqueeze(2),va.unsqueeze(2)],dim=2)\n",
    "        #print(input_tensor.shape)\n",
    "        if not self.batch_first:\n",
    "            # (t, b, c, h, w) -> (b, t, c, h, w)\n",
    "            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)\n",
    "\n",
    "        b, _, _, h, w = input_tensor.size()\n",
    "        # Implement stateful ConvLSTM\n",
    "        if hidden_state is not None:\n",
    "            raise NotImplementedError()\n",
    "        else:\n",
    "            # Since the init is done in forward. Can send image size here\n",
    "            hidden_state = self._init_hidden(batch_size=b,\n",
    "                                             image_size=(h, w))\n",
    "\n",
    "        layer_output_list = []\n",
    "        last_state_list = []\n",
    "\n",
    "        seq_len = input_tensor.size(1)\n",
    "        cur_layer_input = input_tensor\n",
    "\n",
    "        for layer_idx in range(self.num_layers):\n",
    "            h, c = hidden_state[layer_idx]\n",
    "            output_inner = []\n",
    "            for t in range(seq_len):\n",
    "                h, c = self.cell_list[layer_idx](input_tensor=cur_layer_input[:, t, :, :, :],\n",
    "                                                 cur_state=[h, c])\n",
    "                output_inner.append(h)\n",
    "\n",
    "            layer_output = torch.stack(output_inner, dim=1)\n",
    "            cur_layer_input = layer_output\n",
    "\n",
    "            layer_output_list.append(layer_output)\n",
    "            last_state_list.append([h, c])\n",
    "        if not self.return_all_layers:\n",
    "            layer_output_list = layer_output_list[-1:]\n",
    "            last_state_list = last_state_list[-1:]\n",
    "\n",
    "        # return layer_output_list, last_state_list\n",
    "        \n",
    "        #pre=np.array(layer_output_list)[0]\n",
    "        #return pre[0,0,-1,-1,0] \n",
    "        output = layer_output_list[0][:,-1,:,:,:]   #([32, -1, 16, 24, 72])\n",
    "        output = output.view(-1,8*24*72)\n",
    "        output = self.liner(output)\n",
    "        return output\n",
    "\n",
    "    def _init_hidden(self, batch_size, image_size):\n",
    "        init_states = []\n",
    "        for i in range(self.num_layers):\n",
    "            init_states.append(self.cell_list[i].init_hidden(batch_size, image_size))\n",
    "        return init_states\n",
    "\n",
    "    @staticmethod\n",
    "    def _check_kernel_size_consistency(kernel_size):\n",
    "        if not (isinstance(kernel_size, tuple) or\n",
    "                (isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):\n",
    "            raise ValueError('`kernel_size` must be tuple or list of tuples')\n",
    "\n",
    "    @staticmethod\n",
    "    def _extend_for_multilayer(param, num_layers):\n",
    "        if not isinstance(param, list):\n",
    "            param = [param] * num_layers\n",
    "        return param\n",
    "    \n",
    "\n",
    "\n",
    "def coreff(x, y):\n",
    "    x_mean = np.mean(x)\n",
    "    y_mean = np.mean(y)\n",
    "    c1 = sum((x - x_mean) * (y - y_mean))\n",
    "    c2 = sum((x - x_mean)**2) * sum((y - y_mean)**2)\n",
    "    return c1/np.sqrt(c2)\n",
    "\n",
    "def rmse(preds, y):\n",
    "    return np.sqrt(sum((preds - y)**2)/preds.shape[0])\n",
    "\n",
    "def eval_score(preds, label):\n",
    "    # preds = preds.cpu().detach().numpy().squeeze()\n",
    "    # label = label.cpu().detach().numpy().squeeze()\n",
    "    acskill = 0\n",
    "    RMSE = 0\n",
    "    a = 0\n",
    "    a = [1.5]*4 + [2]*7 + [3]*7 + [4]*6\n",
    "    for i in range(24):\n",
    "        RMSE += rmse(label[:, i], preds[:, i])\n",
    "        cor = coreff(label[:, i], preds[:, i])\n",
    "    \n",
    "        acskill += a[i] * np.log(i+1) * cor\n",
    "    return 2/3 * acskill - RMSE\n",
    "\n",
    "\n",
    "\n",
    "device = 'cpu'\n",
    "def train():\n",
    "    fit_params = {\n",
    "    'n_epochs' : 10,\n",
    "    'learning_rate' : 0.0001,\n",
    "    'batch_size' : 100,}\n",
    "\n",
    "    set_seed()\n",
    "    best_sco = -1000\n",
    "    train_dataset, valid_dataset = load_data2()      \n",
    "    train_loader = DataLoader(train_dataset, batch_size=fit_params['batch_size'])\n",
    "    valid_loader = DataLoader(valid_dataset, batch_size=fit_params['batch_size'])\n",
    "    best_sco = -1000\n",
    "    #model = simpleSpatailTimeNN()\n",
    "    model = ConvLSTM(4, 8, (3,3), 1, True, True, False)\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=fit_params['learning_rate'])\n",
    "    loss_fn = nn.MSELoss()   \n",
    "    \n",
    "    model.to(device)\n",
    "    loss_fn.to(device)\n",
    "\n",
    "    for i in range(fit_params['n_epochs']):\n",
    "        model.train()\n",
    "        for step, ((sst, t300, ua, va), label) in enumerate(train_loader): \n",
    "            #print(sst.shape)\n",
    "            sst = sst.to(device).float()\n",
    "            t300 = t300.to(device).float()\n",
    "            ua = ua.to(device).float()\n",
    "            va = va.to(device).float()\n",
    "            optimizer.zero_grad()\n",
    "            label = label.to(device).float()\n",
    "            #print(sst.shape)\n",
    "            preds = model(sst,t300,ua,va)\n",
    "            loss = loss_fn(preds, label)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            #print('Step: {}, Train Loss: {}'.format(step, loss))\n",
    "            \n",
    "        model.eval()\n",
    "        y_true, y_pred = [], []\n",
    "        for step, ((sst, t300, ua, va), label) in enumerate(train_loader):\n",
    "            sst = sst.to(device).float()\n",
    "            t300 = t300.to(device).float()\n",
    "            ua = ua.to(device).float()\n",
    "            va = va.to(device).float()\n",
    "            label = label.to(device).float()\n",
    "            preds = model(sst,t300,ua,va)\n",
    "\n",
    "            y_pred.append(preds)\n",
    "            y_true.append(label)\n",
    "\n",
    "        y_true = torch.cat(y_true, axis=0)\n",
    "        y_pred = torch.cat(y_pred, axis=0)\n",
    "        score = eval_score(y_true.cpu().detach().numpy(), y_pred.cpu().detach().numpy())\n",
    "        ls = loss_fn(y_pred,y_true).cpu().detach().numpy()\n",
    "        print(\"Epoch: {} Train acc: {}  Train loss: {}\".format(i+1,round(score,2),ls))\n",
    "        \n",
    "#         model.eval()\n",
    "        y_true, y_pred = [], []\n",
    "        for step, ((sst, t300, ua, va), label) in enumerate(valid_loader):\n",
    "            sst = sst.to(device).float()\n",
    "            t300 = t300.to(device).float()\n",
    "            ua = ua.to(device).float()\n",
    "            va = va.to(device).float()\n",
    "            label = label.to(device).float()\n",
    "            preds = model(sst,t300,ua,va)\n",
    "\n",
    "            y_pred.append(preds)\n",
    "            y_true.append(label)\n",
    "\n",
    "        y_true = torch.cat(y_true, axis=0)\n",
    "        y_pred = torch.cat(y_pred, axis=0)\n",
    "        score = eval_score(y_true.cpu().detach().numpy(), y_pred.cpu().detach().numpy())\n",
    "        ls = loss_fn(y_pred,y_true).cpu().detach().numpy()\n",
    "        print(\"Epoch: {} Test acc: {}   Test loss: {}\".format(i+1,round(score,2),ls))\n",
    "        \n",
    "        if score>best_sco:\n",
    "            best_sco = score\n",
    "            PATH = './user_data/retrainref.pt'\n",
    "            torch.save(model.state_dict(), PATH)\n",
    "            #torch.save(model, './user_data/ref.pkl')\n",
    "            print('Model saved successfully')\n",
    "\n",
    "\n",
    "\n",
    "def load_data2_SODA(ratio = 0.8):  \n",
    "    # CMIP data    \n",
    "    #PATH = './tcdata/enso_round1_train_20210201/'\n",
    "    #PATH = 'E:/enso_round1_train_20210201/'\n",
    "    PATH = './data/'\n",
    "    train2 = xr.open_dataset(PATH+'SODA_train.nc')\n",
    "    label2 = xr.open_dataset(PATH+'SODA_label.nc')  \n",
    "    train_sst2 = train2['sst'][:, :12].values  # (4645, 12, 24, 72)\n",
    "    train_t3002 = train2['t300'][:, :12].values\n",
    "    train_ua2 = train2['ua'][:, :12].values\n",
    "    train_va2 = train2['va'][:, :12].values\n",
    "    train_label2 = label2['nino'][:, 12:36].values\n",
    "    \n",
    "    m = np.zeros(train_ua2.shape[0])\n",
    "    for i in range(train_ua2.shape[0]):\n",
    "        if np.sum(np.isnan(train_ua2[i])) != 0:\n",
    "            m[i] = np.sum(np.isnan(train_ua2[i]))\n",
    "\n",
    "    #将有缺失值的数据 剔除，选择缺失值为0的数据\n",
    "    train_sst2 = train_sst2[np.where(m==0)[0]]\n",
    "    train_t3002 = train_t3002[np.where(m==0)[0]]\n",
    "    train_ua2 = train_ua2[np.where(m==0)[0]]\n",
    "    train_va2 = train_va2[np.where(m==0)[0]]\n",
    "    train_label2 = train_label2[np.where(m==0)[0]]\n",
    "    \n",
    "    \n",
    "    #测试集比例变动对其结果影响很大！\n",
    "    N = int(len(train_sst2)* ratio)\n",
    "    print('Train samples: {}, Valid samples: {}'.format(len(train_label2[:N]), len(train_label2[N:])))\n",
    "    \n",
    "    dict_train = {\n",
    "        'sst':     train_sst2[2:N],\n",
    "        't300':   train_t3002[2:N],\n",
    "        'ua':       train_ua2[2:N],\n",
    "        'va':       train_va2[2:N],\n",
    "        'label': train_label2[2:N]}\n",
    "    dict_valid = {\n",
    "        'sst':     train_sst2[N:],\n",
    "        't300':   train_t3002[N:],\n",
    "        'ua':       train_ua2[N:],\n",
    "        'va':       train_va2[N:],\n",
    "        'label': train_label2[N:]}\n",
    "    \n",
    "    train_dataset = EarthDataSet(dict_train)\n",
    "    valid_dataset = EarthDataSet(dict_valid)\n",
    "    return train_dataset, valid_dataset\n",
    "\n",
    "\n",
    "\n",
    "def retrain():\n",
    "    set_seed()\n",
    "    best_sco = -1000\n",
    "    fit_params = {\n",
    "        'n_epochs' : 14,\n",
    "        'learning_rate' : 0.0001,\n",
    "        'batch_size' : 100,\n",
    "    }\n",
    "    \n",
    "    train_dataset, valid_dataset = load_data2_SODA()      \n",
    "    train_loader = DataLoader(train_dataset, batch_size=fit_params['batch_size'])\n",
    "    valid_loader = DataLoader(valid_dataset, batch_size=fit_params['batch_size'])\n",
    "    \n",
    "    # Specify a path\n",
    "    PATH = './user_data/retrainref.pt'\n",
    "    # Load\n",
    "#     model = simpleSpatailTimeNN()\n",
    "    model = ConvLSTM(4, 8, (3,3), 1, True, True, False)\n",
    "    model.load_state_dict(torch.load(PATH))\n",
    "    device = 'cpu'\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=fit_params['learning_rate'])\n",
    "    loss_fn = nn.MSELoss()   \n",
    "    \n",
    "    model.to(device)\n",
    "    loss_fn.to(device)\n",
    "\n",
    "    for i in range(fit_params['n_epochs']):\n",
    "        model.train()\n",
    "        for step, ((sst, t300, ua, va), label) in enumerate(train_loader): \n",
    "            #print(sst.shape)\n",
    "            sst = sst.to(device).float()\n",
    "            t300 = t300.to(device).float()\n",
    "            ua = ua.to(device).float()\n",
    "            va = va.to(device).float()\n",
    "            optimizer.zero_grad()\n",
    "            label = label.to(device).float()\n",
    "            #print(sst.shape)\n",
    "            preds = model(sst,t300,ua,va)\n",
    "            loss = loss_fn(preds, label)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            #print('Step: {}, Train Loss: {}'.format(step, loss))\n",
    "            \n",
    "        model.eval()\n",
    "        y_true, y_pred = [], []\n",
    "        for step, ((sst, t300, ua, va), label) in enumerate(train_loader):\n",
    "            sst = sst.to(device).float()\n",
    "            t300 = t300.to(device).float()\n",
    "            ua = ua.to(device).float()\n",
    "            va = va.to(device).float()\n",
    "            label = label.to(device).float()\n",
    "            preds = model(sst,t300,ua,va)\n",
    "\n",
    "            y_pred.append(preds)\n",
    "            y_true.append(label)\n",
    "\n",
    "        y_true = torch.cat(y_true, axis=0)\n",
    "        y_pred = torch.cat(y_pred, axis=0)\n",
    "        score = eval_score(y_true.cpu().detach().numpy(), y_pred.cpu().detach().numpy())\n",
    "        ls = loss_fn(y_pred,y_true).cpu().detach().numpy()\n",
    "        print(\"Epoch: {} Train acc: {}  Train loss: {}\".format(i+1,round(score,2),ls))\n",
    "        \n",
    "#         model.eval()\n",
    "        y_true, y_pred = [], []\n",
    "        for step, ((sst, t300, ua, va), label) in enumerate(valid_loader):\n",
    "            sst = sst.to(device).float()\n",
    "            t300 = t300.to(device).float()\n",
    "            ua = ua.to(device).float()\n",
    "            va = va.to(device).float()\n",
    "            label = label.to(device).float()\n",
    "            preds = model(sst,t300,ua,va)\n",
    "\n",
    "            y_pred.append(preds)\n",
    "            y_true.append(label)\n",
    "\n",
    "        y_true = torch.cat(y_true, axis=0)\n",
    "        y_pred = torch.cat(y_pred, axis=0)\n",
    "        score = eval_score(y_true.cpu().detach().numpy(), y_pred.cpu().detach().numpy())\n",
    "        ls = loss_fn(y_pred,y_true).cpu().detach().numpy()\n",
    "        print(\"Epoch: {} Test acc: {}   Test loss: {}\".format(i+1,round(score,2),ls))\n",
    "        \n",
    "#         if score>best_sco:\n",
    "#             best_sco = score\n",
    "        PATH = './user_data/ref.pt'\n",
    "        torch.save(model.state_dict(), PATH)\n",
    "        #torch.save(model, './user_data/ref.pkl')\n",
    "        print('Model saved successfully')\n",
    "\n",
    "\n",
    "def submit():\n",
    "    # Specify a path\n",
    "    PATH = './user_data/ref.pt'\n",
    "    # Load\n",
    "    #model = simpleSpatailTimeNN()\n",
    "    model = ConvLSTM(4, 8, (3,3), 1, True, True, False)\n",
    "    model.load_state_dict(torch.load(PATH,map_location='cpu'))\n",
    "    model.eval()\n",
    "\n",
    "    #test_path = './tcdata/enso_final_test_data_B/'\n",
    "    test_path = './tcdata/enso_round1_test_20210201/'\n",
    "    #device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    device = 'cpu'\n",
    "    ### 1. 测试数据读取\n",
    "    files = os.listdir(test_path)\n",
    "    test_feas_dict = {}\n",
    "    for file in files:\n",
    "        test_feas_dict[file] = np.load(test_path + file)\n",
    "\n",
    "    ### 2. 结果预测\n",
    "    test_predicts_dict = {}\n",
    "    for file_name, val in test_feas_dict.items():\n",
    "        sst = val[:, :, :, 0]\n",
    "        t300 = val[:, :, :, 1]\n",
    "        ua = val[:, :, :, 2]\n",
    "        va = val[:, :, :, 3]\n",
    "        sst = torch.tensor(sst).float().reshape(-1, 12, 24, 72)\n",
    "        t300 = torch.tensor(t300).float().reshape(-1, 12, 24, 72)\n",
    "        ua = torch.tensor(ua).float().reshape(-1, 12, 24, 72)\n",
    "        va = torch.tensor(va).float().reshape(-1, 12, 24, 72)\n",
    "        test_predicts_dict[file_name] = model(sst, t300, ua, va).reshape(-1, )\n",
    "        print(len(test_predicts_dict))\n",
    "    ### 3.存储预测结果\n",
    "    if not os.path.exists('./result'):\n",
    "        os.makedirs('./result')\n",
    "    for file_name, val in test_predicts_dict.items():\n",
    "        np.save('./result/' + file_name, val.detach().numpy())\n",
    "\n",
    "        # 打包目录为zip文件（未压缩）\n",
    "\n",
    "\n",
    "    def make_zip(source_dir='./result/', output_filename='result.zip'):\n",
    "        zipf = zipfile.ZipFile(output_filename, 'w')\n",
    "        pre_len = len(os.path.dirname(source_dir))\n",
    "        source_dirs = os.walk(source_dir)\n",
    "        print(source_dirs)\n",
    "        for parent, dirnames, filenames in source_dirs:\n",
    "            print(parent, dirnames)\n",
    "            for filename in filenames:\n",
    "                if '.npy' not in filename:\n",
    "                    continue\n",
    "                pathfile = os.path.join(parent, filename)\n",
    "                arcname = pathfile[pre_len:].strip(os.path.sep)  # 相对路径\n",
    "                zipf.write(pathfile, arcname)\n",
    "        zipf.close()\n",
    "\n",
    "\n",
    "    make_zip()\n",
    "    print('==========================end==========================')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train samples: 80, Valid samples: 20\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "Error(s) in loading state_dict for ConvLSTM:\n\tMissing key(s) in state_dict: \"cell_list.0.conv.weight\", \"cell_list.0.conv.bias\", \"liner.weight\", \"liner.bias\". \n\tUnexpected key(s) in state_dict: \"conv1.weight\", \"conv1.bias\", \"batch_norm.running_mean\", \"batch_norm.running_var\", \"batch_norm.num_batches_tracked\", \"lstm.weight_ih_l0\", \"lstm.weight_hh_l0\", \"lstm.bias_ih_l0\", \"lstm.bias_hh_l0\", \"lstm.weight_ih_l0_reverse\", \"lstm.weight_hh_l0_reverse\", \"lstm.bias_ih_l0_reverse\", \"lstm.bias_hh_l0_reverse\", \"lstm.weight_ih_l1\", \"lstm.weight_hh_l1\", \"lstm.bias_ih_l1\", \"lstm.bias_hh_l1\", \"lstm.weight_ih_l1_reverse\", \"lstm.weight_hh_l1_reverse\", \"lstm.bias_ih_l1_reverse\", \"lstm.bias_hh_l1_reverse\", \"linear.weight\", \"linear.bias\". ",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-10-f8ce955600b2>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;31m# train()\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m \u001b[0mretrain\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      3\u001b[0m \u001b[1;31m#submit()\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-9-d85a468f3c21>\u001b[0m in \u001b[0;36mretrain\u001b[1;34m()\u001b[0m\n\u001b[0;32m    448\u001b[0m \u001b[1;31m#     model = simpleSpatailTimeNN()\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    449\u001b[0m     \u001b[0mmodel\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mConvLSTM\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m4\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m8\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m(\u001b[0m\u001b[1;36m3\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;36m3\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;32mFalse\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 450\u001b[1;33m     \u001b[0mmodel\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mload\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mPATH\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    451\u001b[0m     \u001b[0mdevice\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;34m'cpu'\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    452\u001b[0m     \u001b[0moptimizer\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0moptim\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mAdam\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlr\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mfit_params\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m'learning_rate'\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mC:\\software\\Anaconda3\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36mload_state_dict\u001b[1;34m(self, state_dict, strict)\u001b[0m\n\u001b[0;32m   1050\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0merror_msgs\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m>\u001b[0m \u001b[1;36m0\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1051\u001b[0m             raise RuntimeError('Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n\u001b[1;32m-> 1052\u001b[1;33m                                self.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n\u001b[0m\u001b[0;32m   1053\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0m_IncompatibleKeys\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmissing_keys\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0munexpected_keys\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1054\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: Error(s) in loading state_dict for ConvLSTM:\n\tMissing key(s) in state_dict: \"cell_list.0.conv.weight\", \"cell_list.0.conv.bias\", \"liner.weight\", \"liner.bias\". \n\tUnexpected key(s) in state_dict: \"conv1.weight\", \"conv1.bias\", \"batch_norm.running_mean\", \"batch_norm.running_var\", \"batch_norm.num_batches_tracked\", \"lstm.weight_ih_l0\", \"lstm.weight_hh_l0\", \"lstm.bias_ih_l0\", \"lstm.bias_hh_l0\", \"lstm.weight_ih_l0_reverse\", \"lstm.weight_hh_l0_reverse\", \"lstm.bias_ih_l0_reverse\", \"lstm.bias_hh_l0_reverse\", \"lstm.weight_ih_l1\", \"lstm.weight_hh_l1\", \"lstm.bias_ih_l1\", \"lstm.bias_hh_l1\", \"lstm.weight_ih_l1_reverse\", \"lstm.weight_hh_l1_reverse\", \"lstm.bias_ih_l1_reverse\", \"lstm.bias_hh_l1_reverse\", \"linear.weight\", \"linear.bias\". "
     ]
    }
   ],
   "source": [
    "train()\n",
    "retrain()\n",
    "submit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
